#!/usr/bin/python
'''
Takes a tab-delimited file indicating bad regions in contigs and removes edges where 
one of the edge ends maps in the bad region. 
The file format is as follows:
1: something
2: contig
3: bad_start
4: bad_stop
'''
import sys
import os
import numpy as np
import overlap

# parameters.
edge_file = sys.argv[1]
node_file = sys.argv[2]
filter_file = sys.argv[3]
out_file = sys.argv[4]

#####
# read in the filter file.
print "reading filter."
fin = open(filter_file, "rb")
target_lines = fin.readlines()
fin.close()



print "building dictionary."
j = 0
idx = 0
ref_to_idx = {}
targets = np.zeros(len(target_lines), dtype=overlap.entry_dt)
if filter_file.count("gff") > 0:
	# from gff file.
	for i in range(len(target_lines)):
		# skips and token.
		if target_lines[i][0] == "#": continue
		tmp = target_lines[i].strip().split("\t")
		
		# assign idx.
		if tmp[0] not in ref_to_idx:
			ref_to_idx[tmp[0]] = idx
			idx += 1
	
		# add to array.
		targets[j]['type'] = True
		targets[j]['ref'] = ref_to_idx[tmp[0]]
		targets[j]['start'] = int(tmp[3])
		targets[j]['stop'] = int(tmp[4])
		targets[j]['info'] = i
		j += 1
else:
	# from txt file.
	for i in range(len(target_lines)):
		# tokenize.
		tmp = target_lines[i].strip().split("\t")

		# assign idx.
		if tmp[1] not in ref_to_idx:
			ref_to_idx[tmp[1]] = idx
			idx += 1
		
		# add to array.
		targets[j]['type'] = True
		targets[j]['ref'] = ref_to_idx[tmp[1]]
		targets[j]['start'] = int(tmp[2])
		targets[j]['stop'] = int(tmp[3])
		targets[j]['info'] = i
		j += 1


# read in node file.
print "reading node file."
fin = open(node_file, "rb")
lines = fin.readlines()
fin.close()

print "parsing node file."
nodes = {}
for line in lines:
	tmp = line.strip().split("\t")
	nodes[tmp[0]] = tmp[1]


# read in edge file.
print "reading in edges."
fin = open(edge_file, "rb")
edge_lines = fin.readlines()
fin.close()

# build array of starts and filters for each contig.
print "building edge structures.", len(edge_lines)

j = 0
query_to_idx = {}
queries = np.zeros(len(edge_lines) * 2, dtype=overlap.entry_dt)
for i in range(len(edge_lines)):
	# tokenize.
	tmp = edge_lines[i].strip().split("\t")	
	ctg1 = nodes[tmp[0]]
	ctg2 = nodes[tmp[1]]
	pos1 = int(tmp[4])
	pos2 = int(tmp[5])

	# assign idx.
	if ctg1 not in ref_to_idx:
		ref_to_idx[ctg1] = idx
		idx += 1
	if ctg2 not in ref_to_idx:
		ref_to_idx[ctg2] = idx
		idx += 1
	
	# add to array.
	queries[j]['type'] = False
	queries[j]['ref'] = ref_to_idx[ctg1]
	queries[j]['start'] = pos1
	queries[j]['stop'] = pos1 + 50
	queries[j]['info'] = i
	j += 1
	
	queries[j]['type'] = False
	queries[j]['ref'] = ref_to_idx[ctg2]
	queries[j]['start'] = pos2
	queries[j]['stop'] = pos2 + 50
	queries[j]['info'] = i
	j += 1

# Create overlap object.
num_ref = idx + 1
olap = overlap.Overlap(targets, queries, num_ref, verbose=True)

# Find weak overlap.
weako = olap.weak_overlap()

# Flag overlapping edges.
print "flagging overlaps."
j = 0
bad_edges = np.zeros(len(edge_lines), dtype=np.bool)
for i in range(len(weako)):
	# get target, query info.
	#tentry = targets[weako[i][0]]
	qentry = queries[weako[i][1]]
	
	# flag as overlapping.
	bad_edges[qentry['info']] = True

# print unflagged edged.
print "writing good edges."
fout = open(out_file, "wb")
for i in range(len(edge_lines)):
	if bad_edges[i] == False:
		fout.write(edge_lines[i])
fout.close()
